
/*

Arquivo de dados obtido em https://github.com/seade-R/dados-covid-sp

Data inicial = 15/mar/2020 (original = 62 pessoas; filtrado = 118)

Dados filtrados usando Kolmogorov-Zubenko, m=7 pontos, k=3 iterações.

*/

#include <stdio.h>
#include <stdlib.h>

#include <iostream>
#include <fstream>
#include <string>
#include <set>
#include <vector>
#include <random>
#include <algorithm>
#include <map>
#include <cmath>
#include <limits>
#include <iomanip>

//---------------------------------------------------------------------------------

// Auxiliary code 

#include "InverseDistribution.hpp"
#include "csv.hpp"
#include "pid.hpp"
#include "PonteSimulador.c"
#include "arg.h"

//-----------------------------------------------------------------------------

#define M_VERIFICATION_TEST 1

#define Clamp(a,b,c) (a<b)? b : ((a>c) ? c : a)



//-----------------------------------------------------------------------------

#define StopProgramIf(x,str) if(x){std::cout << str << ' ' << __FILE__ << ' ' << __LINE__ << '\n';exit(0);};

//-----------------------------------------------------------------------------

// Core Data-Structures

std::vector<std::vector<int>> Graph;         // for each person, vector linking other persons.
std::vector<bool>             Infected;      //  "    "    "   , state Infected=true.
std::vector<int>              Age;           //
std::vector<int>              DayStartContagion;   //

std::set<int>                 IdxInfected;   // speed-up data structure: list of infected people.
std::map<int, int>            InfectCount;   // number of people infected by one source.


#define MAXDAYS 500
int ReferenceTrajectory[MAXDAYS];
int ReferenceTrajectoryTotal[MAXDAYS];

int NumberInfectedToday = 0;
int NumberInfectedTotal = 0;
int NumberRemovedTotal = 0;

// Maximum number of persons Infected by the Same single Carrier.
int Misc = 0;

//-----------------------------------------------------------------------------

//  Random generation department

std::normal_distribution<double>       RndNormal (0.0, 1.0);
std::uniform_real_distribution<double> RndUniform(0.0, 1.0);
std::uniform_int_distribution<long>    RndInt(0, 1000000000);
std::mt19937 Generator;


CInvDistribution AgePyramid (&Generator);

CProbTable       InfectProbability;

#define Probability(p) (p>RndUniform(Generator))
//-----------------------------------------------------------------------------


std::ofstream arqcontrole;
std::ofstream arqsaida;
std::ofstream arqtree;

std::vector<std::vector<double>> Control;


//XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

//       Core Modeling Alternatives


#include "NetworkModels.cpp"
#include "InfectionModels.cpp"


//XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX







//---------------------------------------------------------
// Randomly infects #i persons
//---------------------------------------------------------
void Setup_Day_Zero (int population, int n) {

   Infected.clear  ();
   Infected.resize (population, false);
   std::fill       (Infected.begin(),Infected.end(),0);

   DayStartContagion.clear ();
   DayStartContagion.resize (population, 0);
   std::fill       (DayStartContagion.begin(),DayStartContagion.end(),-1);

   IdxInfected.clear ();

   InfectCount.clear ();

   int k;

   while (n--) {

      k = RndUniform(Generator) * population;

      if (Infected[k]) continue;

      Infected[k] = true;

      IdxInfected.insert (k);

      DayStartContagion[k] = 0;
   }

}


//=========================================================
// ./a.out   (otimizando)
// ./a.out 1 2 3  (testando otimização)
int main (int argc, char *argv[]) {

   int
      population
     ,num_neighbors
     ,num_infected_day_1
     ,num_days
     ,period_incubation
     ,period_contagious
     ,day_infection_stopped = 0
     ,random_seed
     ;

   double 
      infection_rate
     ,graph_irregular
     ,graph_deg_range
     ,pid_gain = -1.
     ,pid_p, pid_i, pid_d
     ;



   //---------------------------------------------------
   std::string 
       graph_filename
      ,graph_dot_filename
      ,tree_dot_filename
      ,pid_input_filename
      ,pid_output_filename
      ,data_input_filename
      ,quarentine_filename
      ,cases_filename
      ,results_filename
      ;




//==========================================================================
//==========================================================================
//==========================================================================
//==========================================================================


   // There are two points where the random seed is set.
   // The first one is here, to make sure that the code is in control of
   //    the pseudo-random sequence (i.e., making execution deterministic).
   // This seed control generation of population data, and 
   //    generation of the small-world graph that is stored in disk.
   // When the simulator is executed in optimization mode,
   //    the graph is loaded from disk: at that point,
   //    the random seed is set to another value to guarantee
   //    deterministic execution of the simulation.
   // That's a simpler say to circumvent storing/loading generator state,
   //    and it is unlikely harmful.
   Generator.seed (123);

   population         = 12000000;
   num_infected_day_1 = 30;

   num_neighbors   = 50;
   graph_deg_range = 5;
   graph_irregular = 0.1;

   num_days        = 200;  // Até 11 de Setembro.

   //    https://doi.org/10.1038/s41586-020-2456-9
   //    https://doi.org/10.1101/2020.09.04.20188516
   period_contagious    = 5;
   period_incubation    = 3;
    
   graph_filename       = "";//"/home/koscianski/Downloads/Graph12M.bin";

   data_input_filename   = "SP.csv";

   graph_dot_filename    = ""; 
   tree_dot_filename     = "";"Tree50n.dot";
   cases_filename        = "";"Cases50n.csv";
   results_filename      = "";//"sim50nt.csv";
   pid_input_filename    = "";//ctrl12-50.csv";
   pid_output_filename   = "";//"ctrl12-50.csv";

   char *str;

   if (NULL != (str = arg2str (argv, "pidin")))    pid_input_filename   = std::string(str);
   if (NULL != (str = arg2str (argv, "pidout")))   pid_output_filename  = std::string(str);

   if (NULL != (str = arg2str (argv, "graph")))    graph_filename       = std::string(str);

   if (NULL != (str = arg2str (argv, "results")))  results_filename     = std::string(str);
   if (NULL != (str = arg2str (argv, "cases")))    cases_filename       = std::string(str);
   if (NULL != (str = arg2str (argv, "tree")))     tree_dot_filename    = std::string(str);

   if (!arg2int (argv, "randseed", random_seed))
      random_seed = 321;

   if (   (NULL == arg2str (argv, "OPT")) 
       && (NULL == arg2str (argv, "GENGRAPH"))
       && (NULL == arg2str (argv, "PIDPARM"))
       && (NULL == arg2str (argv, "PIDHIST"))
      ) {
      std::cout << "\n\nEither call\n"
                << "sim OPT      to optimize\n"
                << "sim GENGRAPH to generate graph\n"
                << "sim PIDPARM  to run with given PID parameters\n"
                << "sim PIDHIST  to run with given PID history from file\n";

      exit (0);

   }


   bool flag_gengraph    = arg2str (argv, "GENGRAPH");
   bool flag_optimizing  = arg2str (argv, "OPT");
   bool flag_pidhist     = arg2str (argv, "PIDHIST");
   bool flag_pidparm     = arg2str (argv, "PIDPARM");


   if (!flag_gengraph && ("" == graph_filename)) {

      std::cout << "\n\nEither call\n"
                << "sim GENGRAPH graph filename.bin to generate graph   OR\n"
                << "sim other-options with graph filaname.bin to load graph.\n";

      exit (0);

   }


   if (  (flag_optimizing &&  (flag_pidhist || flag_pidparm)) 
       || (!flag_optimizing && !flag_pidhist && !flag_pidparm)
      )  {

      std::cout << "\n\nEither call\n"
                << "sim OPT      to optimize\n"
                << "sim GENGRAPH to generate graph\n"
                << "sim PIDPARM  to run with given PID parameters\n"
                << "sim PIDHIST  to run with given PID history from file\n";

      exit (0);

   }


   if (flag_optimizing &&  ("" == pid_output_filename)) {

      std::cout << "\n\nCall\n"
                << "sim OPT pidout outputPidFile.csv\n";

      exit (0);
   }


   if (arg2str (argv, "GENGRAPH")) {

      if ("" == graph_filename) {
         std::cout << "\nCall ./sim GENGRAPH graph filename.bin\n";
         exit (1);
      }

      std::cout << "\nGenerating the graph...\n";

      WattsStrogatz (Graph, population, num_neighbors, 
                 (int) (graph_deg_range * num_neighbors), graph_irregular);

      std::cout << "... done.\n";
      SaveGraph (graph_filename, Graph);
      exit (0);
   }


   if (!flag_optimizing && flag_pidhist && (pid_input_filename == "")) {

      std::cout << "\n\nCall\n"
                << "sim PIDHIST pidin pidfile.csv graph thegraph.bin\n";

      exit (0);
   }

   if (!flag_optimizing && flag_pidparm) {

      if (!arg2double (argv, "PIDPARM", pid_p)) {
         std::cout << "Expected ./sim PIDPARM 0.1 0.2 0.3 \n";
         exit (0);
      }

      pid_i = std::stod (arg2next (argv));
      pid_d = std::stod (arg2next (argv));

   }

//==========================================================================
//==========================================================================
//==========================================================================
//==========================================================================


   // Read control file (= infection rate). 
   int nrows, ncols;

   // if control should be read from a file.
   if (pid_input_filename != "") {

      Control = (readCSV ((char*) pid_input_filename.c_str(), ',', nrows, ncols, false));

      if (nrows < num_days) num_days = nrows;

      std::cout << "\nControl (" << nrows << " lines) read from [" 
                << pid_input_filename << "]\n";
   }

   // if control should be registered.
   if (pid_output_filename != "") {

      arqcontrole.open (pid_output_filename);

      if (arqcontrole.fail()) { 
		   std::cout << "\nCannot create control output file [" << pid_output_filename << "]\n";
		   exit (0);
      }
   }

   if (tree_dot_filename != "") {

      arqtree.open (tree_dot_filename);

      if (arqtree.fail()) { 
		   std::cout << "\nCannot create file [" << tree_dot_filename << "]\n";
		   exit (0);
      }

      arqtree << "strict graph {\nnode[label=\"\"]\nnode[shape=point]\n";
   }

   if (!flag_optimizing) {

      arqsaida.open (results_filename);

      if (arqsaida.fail()) {
		   std::cout << "\nCannot create file [" << results_filename << "]\n";
		   exit (0);
      }

      arqsaida << "Day,Control,Infected,Reference,Total,Reference_Total\n";
      arqsaida << std::fixed << std::setprecision (std::numeric_limits<double>::digits10 + 2);
   }



   //============================================================================

   //  DEMOGRAPHICS

   // aage/afreq = age pyramid from seade.gov.br, city of São Paulo.
   std::vector<double> aage  {0    , 5    , 10   , 15   , 20   , 25  , 30  , 35   , 40   , 45  , 50   , 55   , 60  , 65  , 70  , 75   , 80   , 85   , 90   };
   std::vector<double> afreq {0.055, 0.061, 0.061, 0.065, 0.076, 0.08, 0.08, 0.082, 0.081, 0.07, 0.065, 0.058, 0.05, 0.04, 0.03, 0.021, 0.014, 0.007, 0.003};

// Estatísticas para o estado de SP: IBGE
//{0   ,5    ,10   ,15   ,20   ,25   ,30   ,35   ,40   ,45   ,50   ,55   ,60   ,65   ,70   ,75   ,80   ,85   ,90    ,95};
//{.067,.0721,.0838,.0833,.0917,.0956,.0897,.0803,.0752,.0694,.0614,.0501,.0227,.0166,.0106,.0151,.0100,.0019,.00167,0.0004};



   // iage/iinfection = distribution of infected by age, source = doi.org/10.1590/0102-311X00149420 
   std::vector<double> iage  {0     ,5      ,10      ,20    ,40    ,60};
   std::vector<double> ifreq {.008  ,.001   ,.005    ,.157  ,.377  ,.452};

   //-------------------------------------------------------------------------------------------------------------

   // Prepare pdf (probability distribution function) with demographic data
   AgePyramid.prepare (afreq, aage);

   // This array will store age of each individual
   Age.resize (population);

   // Assign random age according to the pdf
   for (auto &a : Age)
      a = (unsigned char) AgePyramid.ddraw();

   // Prepare another pdf, this time with infection rates.
   InfectProbability.prepare (ifreq, iage);


   //============================================================================

   // Read data file (= number of infected people). 
   if (data_input_filename != "") {

      std::vector<std::vector<double>> auxiliary = 
           (readCSV ((char*) data_input_filename.c_str(), ',', nrows, ncols, false));

      int t = auxiliary.size();

      if (t < 1) {
         std::cout << "\nTrouble reading file:\n  " << data_input_filename << "\n";
         exit (0);
      }

      t = (t > MAXDAYS) ? MAXDAYS : t;

      while (--t > -1) {
         ReferenceTrajectory[t]      = (int) auxiliary[t][0];
         ReferenceTrajectoryTotal[t] = (int) auxiliary[t][1];
      }

      num_infected_day_1 = ReferenceTrajectoryTotal[0];

   } else {
      std::cout << "\nData file (infection curve) not found [" << data_input_filename << "]\n";
      exit (0);
   }


   //============================================================================
   //
   // Graph Topology
   //
   // doi 10.1287/mnsc.1070.0787   always 10 links, various graphs (Watts-Strogatz,
   //    scale-free network, ring, random...
   //
   // http://dx.doi.org/10.1145/2808797.2808870  average degree 60, pg 1329
   //
   // doi: 10.1109/WSC.2009.5429425 média de 4 vizinhos  between 53 and 56, pg 1010
   //
   //============================================================================

   // Prepare graph (this may take some time)

   if (graph_filename != "") {

      std::cout << "\nLoading graph from [" << graph_filename << "]...\n";

      if (!LoadGraph (graph_filename, Graph)) {
         std::cout << "Failed to open " << graph_filename << "\n";
         exit (1);
      }

      std::cout << "... done.\n";

   }


   //=========================================================================================================

   double error;
   double errorPoissonlikelihood;


   // If Optimizing, this loop will iterate;
   do {

      infection_rate  = 0.01;

      if (flag_optimizing)
         if (!ConectarNomad ()) {
            fprintf (stderr, "Erro de conexão\n");
            exit (-1);
         }

      error = 0;
      errorPoissonlikelihood = 0;

      // This is the second point of the software where the
      //    random seed is set.
      // Read comment on the first point (above).
      Generator.seed (random_seed); // Either 321 or what the user choose.


      // Read PID parameters from Nomad.
      double parameters[3];

      if (flag_optimizing) {

         ReceberNomad (parameters);

         pid_p = parameters[0];
         pid_i = parameters[1];
         pid_d = parameters[2];

         //std::cout << "parms = " << pid_p << ' ' << pid_i << ' ' << pid_d ;

      }

      //---------------------------------------------------------------------

      //===========================================================================================

      // WARNING. The meaning of variable "infect" varies according to the
      //    logic implemented in "Run_One_Day".
      // It is simply not worth the effort to create a completely configurable code;
      //    this program is a workbench itself, as organized and clean as possible,
      //    but due to the nature of the Infection-Model-Logic, I prefer to 
      //    leave loose ends, as pipes that I can reconnect to design each different experiment.

      int clock;

      // For each new infection rate, clean the population and contaminate 
      //   a fell lucky people.
      Setup_Day_Zero (population, num_infected_day_1);

      NumberInfectedTotal = num_infected_day_1;
      NumberRemovedTotal  = 0;

      // Reset PID control.
      PID_Ctrl (NAN, 0, 0, 0);

      // For N days, iterate.
      clock = -1;

      // Welcome to the MAIN SIMULATION LOOP
      while (++clock < num_days) {

         // Record results.
         if (!flag_optimizing) 
            arqsaida         << clock                        
                     << ','  << infection_rate
                     << ','  << NumberInfectedToday
                     << ','  << ReferenceTrajectory[clock]
                     << ','  << NumberInfectedTotal
                     << ','  << ReferenceTrajectoryTotal[clock] 
                     << '\n'; 

         // Record control again in another file, for convenience.
         if (arqcontrole.is_open()) {
            arqcontrole << std::fixed << std::setprecision (std::numeric_limits<double>::digits10 + 2);
	         arqcontrole << infection_rate << '\n';
	         arqcontrole.flush();
         }

         if  (IdxInfected.size() > 0) {

            if (pid_input_filename != "")
                infection_rate = Control[clock][0];

            // Run contagion algorithm for today.
            Model  (clock
                    ,infection_rate 
                    ,period_incubation
                    ,period_contagious);


            // Compute control for tomorrow,
            //    considering what happened today.
            if (pid_input_filename == "") {
               double v = PID_Ctrl (ReferenceTrajectoryTotal[clock] - NumberInfectedTotal
                                          ,pid_p
                                          ,pid_i
                                          ,pid_d);

               infection_rate = 0.001 * v;
               infection_rate = Clamp (infection_rate, 0.005, 0.5);
            }
         } 

         double dif = ReferenceTrajectoryTotal[clock] - NumberInfectedTotal;

         error = error + (dif * dif);
         errorPoissonlikelihood = errorPoissonlikelihood + 
                     ReferenceTrajectory[clock] * log(NumberInfectedToday) - NumberInfectedToday;

      } // while (main simulation loop)


      error = sqrt (error);
      //errorPoissonlikelihood  ** nothing to do **

      // !!! CAREFUL !!!
      // This changes the metric used for optimization.
      // !!! CAREFUL !!!
      error = -errorPoissonlikelihood;

      if (flag_optimizing)
         EnviarNomad (error);

      std::cout << error << ','
                << pid_p << ',' << pid_i << ',' << pid_d << '\n';

   // if optimizing, do it again.
   } while (flag_optimizing);


   // arqsaida was being written to.
   if (!flag_optimizing) {

      arqsaida.close ();

      arqcontrole.close ();

      if (arqtree.is_open()) {
         arqtree << "}\n";
         arqtree.close ();
      }

      if (cases_filename != "") {

         std::ofstream arqcases;
         arqcases.open (cases_filename);

         if (!arqcases.fail()) { 

            arqcases << "Id,Age,Date,Descendents\n";

            for (int id = 0; id < population; ++id) {

               auto it = InfectCount.find (id);

               if (DayStartContagion[id] < 0)
                  continue;

               arqcases << id << ','
                        << Age[id] << ',' 
                        << DayStartContagion[id]-period_incubation << ','
                        << (it->second) << '\n';
            }
            arqcases.close ();
         }
      }
   }

   //===========================================================================================

   // Final Statistics.


   // Average non-Infected and Infected Age

   int c  = 0, sum  = 0, 
       ci = 0, sumi = 0;

   for (int t = 0; t < population; t++) {

      // If you're not infected, your age enters one of the output statistics.
      if (DayStartContagion[t] < 1) {
 
         ++c;
         sum += Age[t];

      } else {

         ++ci;
         sumi += Age[t];
      }
   }

   double ania = (double) sum  / c;
   double aia  = (double) sumi / ci;

   int    max = 0;
   double media = 0;

   for (auto p: InfectCount) {

      media += p.second;

      if (max < p.second) max = p.second;
   }

   media /= InfectCount.size ();

   std::cout << "\n-----------------------------------------"
            << "\nEnd of Simulation."
            << "\nInfected   = " << NumberInfectedTotal
            << "\nAverage R0 = " << media 
            << "\nAverage Non-Infected age = " << ania 
            << "\nAverage Infected age     = " << aia 
            << "\nMaximum descendents from one patient = " << max
            << "\n";


   return 0;
}
